!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief  numerical operations on real-space grid
!> \par History
!>       12.2014 created [Hossein Bani-Hashemian]
!> \author Mohammad Hossein Bani-Hashemian
! **************************************************************************************************
MODULE rs_methods

   USE kinds,                           ONLY: dp
   USE pw_grid_types,                   ONLY: pw_grid_type
   USE pw_methods,                      ONLY: pw_integrate_function,&
                                              pw_scale,&
                                              pw_transfer,&
                                              pw_zero
   USE pw_pool_types,                   ONLY: pw_pool_type
   USE pw_types,                        ONLY: pw_c1d_gs_type,&
                                              pw_r3d_rs_type
   USE realspace_grid_types,            ONLY: realspace_grid_desc_type,&
                                              realspace_grid_type,&
                                              rs_grid_create,&
                                              rs_grid_release,&
                                              rs_grid_zero,&
                                              transfer_pw2rs,&
                                              transfer_rs2pw
#include "../base/base_uses.f90"

   IMPLICIT NONE
   PRIVATE
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'rs_methods'

   PUBLIC derive_fdm_cd3, &
      derive_fdm_cd5, &
      derive_fdm_cd7, &
      setup_grid_axes, &
      pw_mollifier

   REAL(dp), PARAMETER, PRIVATE         :: small_value = 1.0E-8_dp

CONTAINS

! **************************************************************************************************
!> \brief    2nd order finite difference derivative of a function on realspace grid
!> \param f  input function
!> \param df derivative of f
!> \param rs_grid real-space grid
!> \par History:
!>      - Creation (15.11.2013,MK)
!>      - Refactored and moved here from qs_sccs.F (12.2014, Hossein Bani-Hashemian)
!> \author     Matthias Krack (MK)
!> \version    1.0
! **************************************************************************************************
   SUBROUTINE derive_fdm_cd3(f, df, rs_grid)

      TYPE(pw_r3d_rs_type), INTENT(IN)                   :: f
      TYPE(pw_r3d_rs_type), DIMENSION(3), INTENT(INOUT)  :: df
      TYPE(realspace_grid_type), INTENT(IN)              :: rs_grid

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'derive_fdm_cd3'

      INTEGER                                            :: handle, i, j, k
      INTEGER, DIMENSION(3)                              :: lb, ub
      REAL(KIND=dp), DIMENSION(3)                        :: h
      REAL(KIND=dp), DIMENSION(:, :, :), POINTER         :: drdx, drdy, drdz, r
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      TYPE(realspace_grid_type), DIMENSION(3)            :: drs_grid

      CALL timeset(routineN, handle)

      ! Setup
      rs_desc => rs_grid%desc
      CALL transfer_pw2rs(rs_grid, f)
      DO i = 1, 3
         CALL rs_grid_create(drs_grid(i), rs_desc)
         CALL rs_grid_zero(drs_grid(i))
      END DO

      lb(1:3) = rs_grid%lb_real(1:3)
      ub(1:3) = rs_grid%ub_real(1:3)
      r => rs_grid%r
      drdx => drs_grid(1)%r
      drdy => drs_grid(2)%r
      drdz => drs_grid(3)%r

      ! 3-point stencil central differences
      h(1:3) = 2.0_dp*f%pw_grid%dr(1:3)
!$OMP     PARALLEL DO DEFAULT(NONE) &
!$OMP                 PRIVATE(i,j,k) &
!$OMP                 SHARED(drdx,drdy,drdz,h,lb,r,ub)
      DO k = lb(3), ub(3)
         DO j = lb(2), ub(2)
            DO i = lb(1), ub(1)
               drdx(i, j, k) = (r(i + 1, j, k) - r(i - 1, j, k))/h(1)
               drdy(i, j, k) = (r(i, j + 1, k) - r(i, j - 1, k))/h(2)
               drdz(i, j, k) = (r(i, j, k + 1) - r(i, j, k - 1))/h(3)
            END DO
         END DO
      END DO
!$OMP     END PARALLEL DO

      ! Cleanup
      DO i = 1, 3
         CALL transfer_rs2pw(drs_grid(i), df(i))
         CALL rs_grid_release(drs_grid(i))
      END DO

      CALL timestop(handle)

   END SUBROUTINE derive_fdm_cd3

! **************************************************************************************************
!> \brief    4th order finite difference derivative of a function on realspace grid
!> \param f  input function
!> \param df derivative of f
!> \param rs_grid real-space grid
!> \par History:
!>      - Creation (15.11.2013,MK)
!>      - Refactored and moved here from qs_sccs.F (12.2014, Hossein Bani-Hashemian)
!> \author     Matthias Krack (MK)
!> \version    1.0
! **************************************************************************************************
   SUBROUTINE derive_fdm_cd5(f, df, rs_grid)

      TYPE(pw_r3d_rs_type), INTENT(IN)                   :: f
      TYPE(pw_r3d_rs_type), DIMENSION(3), INTENT(INOUT)  :: df
      TYPE(realspace_grid_type), INTENT(IN)              :: rs_grid

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'derive_fdm_cd5'

      INTEGER                                            :: handle, i, j, k
      INTEGER, DIMENSION(3)                              :: lb, ub
      REAL(KIND=dp), DIMENSION(3)                        :: h
      REAL(KIND=dp), DIMENSION(:, :, :), POINTER         :: drdx, drdy, drdz, r
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      TYPE(realspace_grid_type), DIMENSION(3)            :: drs_grid

      CALL timeset(routineN, handle)

      ! Setup
      rs_desc => rs_grid%desc
      CALL transfer_pw2rs(rs_grid, f)
      DO i = 1, 3
         CALL rs_grid_create(drs_grid(i), rs_desc)
         CALL rs_grid_zero(drs_grid(i))
      END DO

      lb(1:3) = rs_grid%lb_real(1:3)
      ub(1:3) = rs_grid%ub_real(1:3)
      r => rs_grid%r
      drdx => drs_grid(1)%r
      drdy => drs_grid(2)%r
      drdz => drs_grid(3)%r

      ! 5-point stencil central differences
      h(1:3) = 12.0_dp*f%pw_grid%dr(1:3)
!$OMP     PARALLEL DO DEFAULT(NONE) &
!$OMP                 PRIVATE(i,j,k) &
!$OMP                 SHARED(drdx,drdy,drdz,h,lb,r,ub)
      DO k = lb(3), ub(3)
         DO j = lb(2), ub(2)
            DO i = lb(1), ub(1)
               drdx(i, j, k) = (r(i - 2, j, k) - r(i + 2, j, k) + 8.0_dp*(r(i + 1, j, k) - r(i - 1, j, k)))/h(1)
               drdy(i, j, k) = (r(i, j - 2, k) - r(i, j + 2, k) + 8.0_dp*(r(i, j + 1, k) - r(i, j - 1, k)))/h(2)
               drdz(i, j, k) = (r(i, j, k - 2) - r(i, j, k + 2) + 8.0_dp*(r(i, j, k + 1) - r(i, j, k - 1)))/h(3)
            END DO
         END DO
      END DO
!$OMP     END PARALLEL DO

      ! Cleanup
      DO i = 1, 3
         CALL transfer_rs2pw(drs_grid(i), df(i))
         CALL rs_grid_release(drs_grid(i))
      END DO

      CALL timestop(handle)

   END SUBROUTINE derive_fdm_cd5

! **************************************************************************************************
!> \brief    6th order finite difference derivative of a function on realspace grid
!> \param f  input function
!> \param df derivative of f
!> \param rs_grid real-space grid
!> \par History:
!>      - Creation (15.11.2013,MK)
!>      - Refactored and moved here from qs_sccs.F (12.2014, Hossein Bani-Hashemian)
!> \author     Matthias Krack (MK)
!> \version    1.0
! **************************************************************************************************
   SUBROUTINE derive_fdm_cd7(f, df, rs_grid)

      TYPE(pw_r3d_rs_type), INTENT(IN)                   :: f
      TYPE(pw_r3d_rs_type), DIMENSION(3), INTENT(INOUT)  :: df
      TYPE(realspace_grid_type), INTENT(IN)              :: rs_grid

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'derive_fdm_cd7'

      INTEGER                                            :: handle, i, j, k
      INTEGER, DIMENSION(3)                              :: lb, ub
      REAL(KIND=dp), DIMENSION(3)                        :: h
      REAL(KIND=dp), DIMENSION(:, :, :), POINTER         :: drdx, drdy, drdz, r
      TYPE(realspace_grid_desc_type), POINTER            :: rs_desc
      TYPE(realspace_grid_type), DIMENSION(3)            :: drs_grid

      CALL timeset(routineN, handle)

      ! Setup
      rs_desc => rs_grid%desc
      CALL transfer_pw2rs(rs_grid, f)
      DO i = 1, 3
         CALL rs_grid_create(drs_grid(i), rs_desc)
         CALL rs_grid_zero(drs_grid(i))
      END DO

      lb(1:3) = rs_grid%lb_real(1:3)
      ub(1:3) = rs_grid%ub_real(1:3)
      r => rs_grid%r
      drdx => drs_grid(1)%r
      drdy => drs_grid(2)%r
      drdz => drs_grid(3)%r

      ! 7-point stencil central differences
      h(1:3) = 60.0_dp*f%pw_grid%dr(1:3)
!$OMP     PARALLEL DO DEFAULT(NONE) &
!$OMP                 PRIVATE(i,j,k) &
!$OMP                 SHARED(drdx,drdy,drdz,h,lb,r,ub)
      DO k = lb(3), ub(3)
         DO j = lb(2), ub(2)
            DO i = lb(1), ub(1)
               drdx(i, j, k) = (r(i + 3, j, k) - r(i - 3, j, k) + 9.0_dp*(r(i - 2, j, k) - r(i + 2, j, k)) + &
                                45.0_dp*(r(i + 1, j, k) - r(i - 1, j, k)))/h(1)
               drdy(i, j, k) = (r(i, j + 3, k) - r(i, j - 3, k) + 9.0_dp*(r(i, j - 2, k) - r(i, j + 2, k)) + &
                                45.0_dp*(r(i, j + 1, k) - r(i, j - 1, k)))/h(2)
               drdz(i, j, k) = (r(i, j, k + 3) - r(i, j, k - 3) + 9.0_dp*(r(i, j, k - 2) - r(i, j, k + 2)) + &
                                45.0_dp*(r(i, j, k + 1) - r(i, j, k - 1)))/h(3)
            END DO
         END DO
      END DO
!$OMP     END PARALLEL DO

      ! Cleanup
      DO i = 1, 3
         CALL transfer_rs2pw(drs_grid(i), df(i))
         CALL rs_grid_release(drs_grid(i))
      END DO

      CALL timestop(handle)

   END SUBROUTINE derive_fdm_cd7

! **************************************************************************************************
!> \brief returns the global axes and the portion of the axes that are local to
!>        the current mpi rank
!> \param pw_grid plane wave grid
!> \param x_glbl x grid vector of the simulation box
!> \param y_glbl y grid vector of the simulation box
!> \param z_glbl z grid vector of the simulation box
!> \param x_locl x grid vector of the simulation box local to this process
!> \param y_locl y grid vector of the simulation box local to this process
!> \param z_locl z grid vector of the simulation box local to this process
!> \par History
!>       07.2014 created [Hossein Bani-Hashemian]
!>       07.2015 moved here from dirichlet_bc_utils.F [Hossein Bani-Hashemian]
!> \author Mohammad Hossein Bani-Hashemian
! **************************************************************************************************
   SUBROUTINE setup_grid_axes(pw_grid, x_glbl, y_glbl, z_glbl, x_locl, y_locl, z_locl)

      TYPE(pw_grid_type), INTENT(IN)                     :: pw_grid
      REAL(dp), ALLOCATABLE, DIMENSION(:), INTENT(OUT)   :: x_glbl, y_glbl, z_glbl, x_locl, y_locl, &
                                                            z_locl

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'setup_grid_axes'

      INTEGER                                            :: glb1, glb2, glb3, gub1, gub2, gub3, &
                                                            handle, i, lb1, lb2, lb3, ub1, ub2, ub3
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: gindx, gindy, gindz, lindx, lindy, lindz
      INTEGER, DIMENSION(2, 3)                           :: bounds, bounds_local
      INTEGER, DIMENSION(3)                              :: npts, npts_local
      REAL(dp), DIMENSION(3)                             :: dr

      CALL timeset(routineN, handle)

      dr = pw_grid%dr
      bounds = pw_grid%bounds
      bounds_local = pw_grid%bounds_local
      npts = pw_grid%npts
      npts_local = pw_grid%npts_local

! local and global lower and upper bounds
      glb1 = bounds(1, 1); gub1 = bounds(2, 1)
      glb2 = bounds(1, 2); gub2 = bounds(2, 2)
      glb3 = bounds(1, 3); gub3 = bounds(2, 3)
      lb1 = bounds_local(1, 1); ub1 = bounds_local(2, 1)
      lb2 = bounds_local(1, 2); ub2 = bounds_local(2, 2)
      lb3 = bounds_local(1, 3); ub3 = bounds_local(2, 3)

      ALLOCATE (lindx(lb1:ub1), lindy(lb2:ub2), lindz(lb3:ub3))
      ALLOCATE (gindx(glb1:gub1), gindy(glb2:gub2), gindz(glb3:gub3))
      ALLOCATE (x_locl(lb1:ub1), y_locl(lb2:ub2), z_locl(lb3:ub3))
      ALLOCATE (x_glbl(glb1:gub1), y_glbl(glb2:gub2), z_glbl(glb3:gub3))

      gindx(:) = (/(i, i=0, npts(1) - 1)/)
      gindy(:) = (/(i, i=0, npts(2) - 1)/)
      gindz(:) = (/(i, i=0, npts(3) - 1)/)
      lindx(:) = (/(i, i=0, npts_local(1) - 1)/)
      lindy(:) = (/(i, i=0, npts_local(2) - 1)/)
      lindz(:) = (/(i, i=0, npts_local(3) - 1)/)

      x_glbl(:) = gindx*dr(1)
      y_glbl(:) = gindy*dr(2)
      z_glbl(:) = gindz*dr(3)

! map [0 .. (npts_local-1)] --> [lb .. ub]
      IF (lb1 .EQ. ub1) THEN
         lindx(:) = lb1
      ELSE
         lindx(:) = lindx(:)*((ub1 - lb1)/(npts_local(1) - 1)) + lb1
      END IF
      IF (lb2 .EQ. ub2) THEN
         lindy(:) = lb2
      ELSE
         lindy(:) = lindy(:)*((ub2 - lb2)/(npts_local(2) - 1)) + lb2
      END IF
      IF (lb3 .EQ. ub3) THEN
         lindz(:) = lb3
      ELSE
         lindz(:) = lindz(:)*((ub3 - lb3)/(npts_local(3) - 1)) + lb3
      END IF

      x_locl(:) = x_glbl(lindx)
      y_locl(:) = y_glbl(lindy)
      z_locl(:) = z_glbl(lindz)

      CALL timestop(handle)

   END SUBROUTINE setup_grid_axes

! **************************************************************************************************
!> \brief convolutes a function with a smoothing kernel K_\zeta
!>                         v * K_\zeta
!> K_\zeta is the standard mollifier defined as:
!>        K_\zeta(x) = \frac{1}{\zeta^3} K(\frac{x}{\zeta})
!> where
!>        K(x) = \kappa \exp (\frac{1}{|x|^2 - 1}),  if |x| <= 1
!>             = 0,                                  otherwise
!> \param pw_pool pool of pw grid
!> \param zeta parameter \zeta defining the width of the mollifier
!> \param x_glbl x grid vector of the simulation box
!> \param y_glbl y grid vector of the simulation box
!> \param z_glbl z grid vector of the simulation box
!> \param pw_in the input function
!> \param pw_out the convoluted function
!> \par History
!>       10.2014 created [Hossein Bani-Hashemian]
!>       07.2015 moved here from ps_implicit_methods.F [Hossein Bani-Hashemian]
!> \author Mohammad Hossein Bani-Hashemian
! **************************************************************************************************
   SUBROUTINE pw_mollifier(pw_pool, zeta, x_glbl, y_glbl, z_glbl, pw_in, pw_out)

      TYPE(pw_pool_type), INTENT(IN), POINTER            :: pw_pool
      REAL(dp), INTENT(IN)                               :: zeta
      REAL(dp), ALLOCATABLE, DIMENSION(:), INTENT(IN)    :: x_glbl, y_glbl, z_glbl
      TYPE(pw_r3d_rs_type), INTENT(IN)                   :: pw_in
      TYPE(pw_r3d_rs_type), INTENT(INOUT)                :: pw_out

      CHARACTER(LEN=*), PARAMETER                        :: routineN = 'pw_mollifier'

      INTEGER                                            :: handle, i, j, k, lb1, lb2, lb3, ub1, &
                                                            ub2, ub3
      INTEGER, DIMENSION(2, 3)                           :: bounds, bounds_local
      REAL(dp)                                           :: normfact, xi, xmax, xmin, yj, ymax, &
                                                            ymin, zk, zmax, zmin
      REAL(dp), DIMENSION(3, 3)                          :: dh
      TYPE(pw_c1d_gs_type)                               :: G_gs, pw_in_gs, pw_out_gs
      TYPE(pw_grid_type), POINTER                        :: pw_grid
      TYPE(pw_r3d_rs_type)                               :: G

      CALL timeset(routineN, handle)

      pw_grid => pw_in%pw_grid
      dh = pw_grid%dh
      bounds_local = pw_grid%bounds_local
      bounds = pw_grid%bounds

      lb1 = bounds_local(1, 1); ub1 = bounds_local(2, 1)
      lb2 = bounds_local(1, 2); ub2 = bounds_local(2, 2)
      lb3 = bounds_local(1, 3); ub3 = bounds_local(2, 3)

      CALL pw_pool%create_pw(G)
      CALL pw_pool%create_pw(G_gs)
      CALL pw_pool%create_pw(pw_in_gs)
      CALL pw_pool%create_pw(pw_out_gs)

      CALL pw_zero(G)
      xmin = x_glbl(bounds(1, 1)); xmax = x_glbl(bounds(2, 1))
      ymin = y_glbl(bounds(1, 2)); ymax = y_glbl(bounds(2, 2))
      zmin = z_glbl(bounds(1, 3)); zmax = z_glbl(bounds(2, 3))

      DO k = lb3, ub3
         DO j = lb2, ub2
            DO i = lb1, ub1
               xi = x_glbl(i); yj = y_glbl(j); zk = z_glbl(k)
               IF (norm2((/(xi - xmin), (yj - ymin), (zk - zmin)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmin), (yj - ymin), (zk - zmin)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmax), (yj - ymax), (zk - zmax)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmax), (yj - ymax), (zk - zmax)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmin), (yj - ymax), (zk - zmax)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmin), (yj - ymax), (zk - zmax)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmax), (yj - ymin), (zk - zmax)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmax), (yj - ymin), (zk - zmax)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmax), (yj - ymax), (zk - zmin)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmax), (yj - ymax), (zk - zmin)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmin), (yj - ymin), (zk - zmax)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmin), (yj - ymin), (zk - zmax)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmin), (yj - ymax), (zk - zmin)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmin), (yj - ymax), (zk - zmin)/)/zeta)**2 - 1))
               ELSE IF (norm2((/(xi - xmax), (yj - ymin), (zk - zmin)/)) .LT. zeta - small_value) THEN
                  G%array(i, j, k) = EXP(1.0_dp/(norm2((/(xi - xmax), (yj - ymin), (zk - zmin)/)/zeta)**2 - 1))
               END IF
            END DO
         END DO
      END DO
      CALL pw_scale(G, (1.0_dp/zeta)**3)
      normfact = pw_integrate_function(G)
      CALL pw_scale(G, 1.0_dp/normfact)

      CALL pw_transfer(G, G_gs)
      CALL pw_transfer(pw_in, pw_in_gs)
      pw_out_gs%array = G_gs%array*pw_in_gs%array
      CALL pw_transfer(pw_out_gs, pw_out)

      ! multiply by the reciprocal of the forward Fourier transform normalization prefactor (here 1/N, by convention)
      CALL pw_scale(pw_out, REAL(pw_grid%ngpts, KIND=dp))
      ! from discrete convolution to continuous convolution
      CALL pw_scale(pw_out, pw_grid%dvol)

      DO k = lb3, ub3
         DO j = lb2, ub2
            DO i = lb1, ub1
               IF (ABS(pw_out%array(i, j, k)) .LE. 1.0E-10_dp) pw_out%array(i, j, k) = 0.0_dp
            END DO
         END DO
      END DO

      CALL pw_pool%give_back_pw(G)
      CALL pw_pool%give_back_pw(G_gs)
      CALL pw_pool%give_back_pw(pw_in_gs)
      CALL pw_pool%give_back_pw(pw_out_gs)
      CALL timestop(handle)

   END SUBROUTINE pw_mollifier

END MODULE rs_methods
